PROBER = "TCPIP0::192.168.0.2::800::SOCKET"
SMU = "GPIB0::17::INSTR"
SWITCH = "GPIB0::22::INSTR"

import pyvisa
rm = pyvisa.ResourceManager()

linkCheck = True
try:
    prober = rm.open_resource(PROBER)
    prober.read_termination = '\r\n'
    print(prober.query("*IDN?"))
except Exception as e:
    print("MPI Prober Station Link Error!")
    print(e)
    prober = None
    linkCheck =  False

try:
    B1500a = rm.open_resource(SMU)
    B1500a.query("*IDN?")
except Exception as e:
    print(f"B1500 SMU Link Error!")
    print(e)
    linkCheck =  False
    B1500a = None
try:
    B2201 = rm.open_resource(SWITCH)
    B2201.query("*IDN?")
except Exception as e:
    print("B2201 Switching Matrix Link Error!")
    print(e)
    linkCheck =  False
    B2201 = None

def contact(prober=prober):
    prober.write("move_chuck_contact")
def sep(prober=prober):
    prober.write("move_chuck_separation")
def nextdie(prober=prober):
    prober.write("map:step_next_die")
def subdie(lineOffset,prober=prober):
    prober.write(f'move_chuck_xy Relative,0,{lineOffset}')

def matrix(stat = 4,b2200=B2201):
    import B2200
    import time
    #1,1234-1234
    #2,1234-5678
    #3,1234-9-12
    #4,openall
    B2200.Load_Path(b2200,stat)
    
#b1500a smu
def transCurve(startV=-2, endV=0.5,Vds=-2,
               chnSource='1',chnDrain='2',chnGate='3',
               ADC='AIT 1,2,1',
               direction='Single Linear',
               reverse=False,
               points=101,
               t_hold=1.0,s_delay=0.0,t_delay=0.0,**kwarg):

    if reverse:
        startV,endV = endV,startV

    B1500a.write('*RST')
    B1500a.write('CN')

    #--------------------IV_Sweep函数测试参数示例（双变量扫描）------------------#
    Unit1=chnGate                    #扫描变量1对应的物理端口为4
    VName1='VGate'                   #扫描变量1的名字为VD，注意电压模式的端口命名为V....,电流模式端口命名为I...
    Direction1=direction    #扫描变量1的方向为Double，并且加电压模式为Linear（更多模式见34）
    Start1=startV                   #扫描变量1的起始电压为-3 V
    end1=endV                     #扫描变量1的终止电压为3 V
    nop1=points                      #扫描变量1的点数目为201（Double的话，实际数据点为201*2）
    Comp1=1e-2                  #扫描变量1的限流值
    range1='Auto'         #扫描变量1的测量Range

    VAR_setup=[Unit1,VName1,Direction1,Start1,end1,nop1,Comp1,range1] #打包扫描变量，如果单变量扫描，不用第二行

    Measurement_setup = []

    Unit=chnDrain                  #测量变量1的物理端口为
    VName='VDrain'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=Vds                 #测量变量1的固定电压为0 V
    Comp=0.1                 #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个

    Unit=chnSource                   #测量变量1的物理端口为
    VName='VSource'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=0                 #测量变量1的固定电压为0 V
    Comp=0.1                  #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个


    #print(Measurement_setup)
    B1500a.timeout = 60000
    B1500a.write('*RST')
    B1500a.write('CN 1,2,3,4')
    import B1500
    res = B1500.IV_sweept(B1500a,VAR_setup,Measurement_setup,ADC=ADC,t_hold=t_hold,s_delay=s_delay,t_delay=t_delay)#AIT 1,2,1 20ms 1,1,125 10ms

    import pandas
    df = pandas.DataFrame(res[1:],columns=['GateV','GateI','DrainV','DrainI','SourceV','SourceI'],dtype=float)
    return df

#b1500a smu
def outputCurve(startV=-2, endV=0.5,Vgs=-2,
               chnSource='1',chnDrain='2',chnGate='3',
               ADC='AIT 1,2,1',
               direction='Single Linear',
               reverse=False,
               points=101,
               t_hold=1.0,s_delay=0.0,t_delay=0.0,**kwarg):

    if reverse:
        startV,endV = endV,startV

    B1500a.write('*RST')
    B1500a.write('CN')

    #--------------------IV_Sweep函数测试参数示例（双变量扫描）------------------#
    Unit1=chnDrain                    #扫描变量1对应的物理端口为4
    VName1='VDrain'                   #扫描变量1的名字为VD，注意电压模式的端口命名为V....,电流模式端口命名为I...
    Direction1=direction    #扫描变量1的方向为Double，并且加电压模式为Linear（更多模式见34）
    Start1=startV                   #扫描变量1的起始电压为-3 V
    end1=endV                     #扫描变量1的终止电压为3 V
    nop1=points                      #扫描变量1的点数目为201（Double的话，实际数据点为201*2）
    Comp1=1e-2                  #扫描变量1的限流值
    range1='Auto'         #扫描变量1的测量Range

    VAR_setup=[Unit1,VName1,Direction1,Start1,end1,nop1,Comp1,range1] #打包扫描变量，如果单变量扫描，不用第二行

    Measurement_setup = []

    Unit=chnGate                  #测量变量1的物理端口为
    VName='VGate'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=Vgs                 #测量变量1的固定电压为0 V
    Comp=0.1                 #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个

    Unit=chnSource                   #测量变量1的物理端口为
    VName='VSource'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=0                 #测量变量1的固定电压为0 V
    Comp=0.1                  #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个


    #print(Measurement_setup)
    B1500a.timeout = 60000
    B1500a.write('*RST')
    B1500a.write('CN 1,2,3,4')
    import B1500
    res = B1500.IV_sweept(B1500a,VAR_setup,Measurement_setup,ADC=ADC,t_hold=t_hold,s_delay=s_delay,t_delay=t_delay)#AIT 1,2,1 20ms 1,1,125 10ms

    import pandas
    df = pandas.DataFrame(res[1:],columns=['GateV','GateI','DrainV','DrainI','SourceV','SourceI'],dtype=float)
    return df

figureSet={
    transCurve:{
        'xlabel':'GateV (V)',
        'ylabel':'SourceI (A)',
        'xcol':"GateV",
        'ycol':"SourceI",
        'islog':True
        },
    outputCurve:{
        'xlabel':'DrainV (V)',
        'ylabel':'DrainI (A)',
        'xcol':"DrainV",
        'ycol':"DrainI",
        'islog':False
    }
           }

#jupyter选择弹窗模式
def autoTest(testMode=transCurve,
             tests=[],
             dies=58,
             lines=3,
             lineOffset=0.342,
             positions=[1,2,3]):
    if not linkCheck:
        return 
    fset = figureSet[testMode]
    import pandas as pd 
    from datetime import datetime
    tStamp = datetime.strftime(datetime.now(),"%Y%m%d_%H%M%S")
    import os
    os.mkdir(f'autoData/{tStamp}')
    xlsFiles = {}
    from itertools import product
    for line,position,test in product(range(1,lines+1),positions,tests):
        fn = f'{test["name"]}-LINE{line}-{str(position)}'
        xlsFiles[fn] = pd.ExcelWriter(f'autoData/{tStamp}/{fn}.xlsx')

    import matplotlib.pyplot as plt
    import numpy as np
    import time

    plt.style.use('default')
    plt.ion() #开启interactive mode 成功的关键函数
    plt.cla()
    plt.xlabel(fset['xlabel'])
    plt.ylabel(fset['ylabel'])
    if fset['islog']:
        plt.yscale('log')
    plt.title('StarAutoProbe')

    import time
    for i in range(dies):
        nextdie()
        for line in range(1,lines+1):
            if line!=1:
                subdie(lineOffset)
            time.sleep(0.3)
            contact()
            time.sleep(0.3)
            for position in positions:
                matrix(position)
                time.sleep(0.2)
                for test in tests:
                    fn = f'{test["name"]}-LINE{line}-{str(position)}'
                    print('\r'+str(i+1)+f'/{dies}:{fn}',end='')
                    try:
                        df = testMode(**test)
                        plt.title("StarAutoProbe: "+str(i+1)+f'/{dies}:{fn}')
                        plt.plot(df['GateV'],np.abs(df['SourceI']))
                        plt.pause(0.1)
                        df.to_excel(xlsFiles[fn],index=False,sheet_name=f"Run{i}")
                    except:
                        print("\r"+str(i+1)+f'/{dies}:{fn}'+"error, skipped!")
                        try:
                            for line,position,test in product(range(1,lines+1),positions,tests):
                                fn = f'{test["name"]}-LINE{line}-{str(position)}'
                                xlsFiles[fn].close()
                            print("Data saved.")
                        except:
                            pass
                matrix(4)
            print('\r'+str(i+1)+'/'+str(dies),end='')
            if not (i == dies-1 and line == lines-1):
                sep()
    plt.title('StarAutoProbe: Complete!')
    plt.pause(0.1)

    for line,position,test in product(range(1,lines+1),positions,tests):
        fn = f'{test["name"]}-LINE{line}-{str(position)}'
        xlsFiles[fn].close() 

def closeIns():
    if prober:
        prober.close()
    if B1500a:
        B1500a.close()
    if B2201:
        B2201.close()
    rm.close()    